
from sklearn.ensemble import RandomForestClassifier
import mlflow
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
import git
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from azureml.core.model import Model
import pickle

class rfImp(object):

    def __init__(self, dataset, X_train,y_train,X_val,y_val,run):
        self.X_train = X_train
        self.y_train = y_train
        self.X_val = X_val
        self.y_val = y_val
        self.dataset = dataset
        self.run = run
        self.rf = RandomForestClassifier(max_depth=10, random_state=0, n_estimators=100)
        self.parameters = {'kernel':('linear', 'rbf'), 'C':[1, 10]}
    
    def rfImpTrn(self):
        '''
        Modelling with random forest method
        '''
        self.rf.fit(self.X_train, self.y_train)
        mlflow.log_param("max_depth", 10)
        mlflow.log_param("random_state", 0)
        mlflow.log_param("n_estimators", 100)

    def rfImpVal(self):
        '''
        Model testing with random forest method.
        We finally get accuracy, fscore, precision and recall values.
        '''
        predicted_rf = self.rf.predict(self.X_val)
        self.acc = accuracy_score(self.y_val, predicted_rf)
        self.fscore = f1_score(self.y_val, predicted_rf, average="macro")
        self.precision = precision_score(self.y_val, predicted_rf, average="macro")
        self.recall = recall_score(self.y_val, predicted_rf, average="macro")
        repo = git.Repo(search_parent_directories=True)
        self.sha = repo.head.object.hexsha
        # Log to AzureML and MLflow
        mlflow.log_param("Test_accuracy", self.acc)
        mlflow.log_param("Precision", self.precision)
        mlflow.log_param("Test_recall", self.recall)
        mlflow.log_param("F-Score", self.fscore)
        mlflow.log_param("Git-sha", self.sha)
        mlflow.sklearn.log_model(self.rf, 'outputs')
    
    def onnxModelSave(self):
        '''
        Convert random forest model into ONNX format file and save to local path.
        '''
        initial_type = [('float_input', FloatTensorType([None, 6]))]
        onx = convert_sklearn(self.rf, initial_types=initial_type)
        with open("outputs/rf.onnx", "wb") as f:
            f.write(onx.SerializeToString())

    def modelRegister(self,workspace):
        '''
        Register Model on AzureML workstation
        '''
        # Register Model on AzureML WS
        model = Model.register(model_path = './outputs/rf.onnx', # this points to a local file 
                            model_name = "random-forest-classifier", # this is the name the model is registered as
                            tags = {'dataset': self.dataset.name, 'version': self.dataset.version, 'hyparameter-C': '1', 'testdata-accuracy': '0.9548'}, 
                            model_framework='pandas==0.23.4',
                            description = "Random forest classifier to predict weather at port of Turku",
                            workspace = workspace)
        print('Name:', model.name)
        print('Version:', model.version)
        # Save the model to the outputs directory for capture
        mlflow.sklearn.log_model(self.rf, 'outputs/rf.onnx')
    
    def modelRegisterSC(self, sc, workspace):
        '''
        Register sc (StandardScaler) for the data standardrization parameters.
        '''
        with open('./outputs/scaler.pkl', 'wb') as scaler_pkl:
            # sc = StandardScaler()
            pickle.dump(sc, scaler_pkl)
        # Register Model on AzureML WS
        scaler = Model.register(model_path = './outputs/scaler.pkl', # this points to a local file 
                            model_name = "scaler", # this is the name the model is registered as
                            tags = {'dataset': self.dataset.name, 'version': self.dataset.version}, 
                            model_framework='pandas==0.23.4',
                            description = "Scaler used for scaling incoming inference data",
                            workspace = workspace)

        print('Name:', scaler.name)
        print('Version:', scaler.version)